{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b50c4af8-fec3-4396-860a-1322089d76cb",
   "metadata": {},
   "source": [
    "# OpenAI Agent with Query Engine Tools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db402a8b-90d6-4e1d-8df6-347c54624f26",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Build Query Engine Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "02160804-64a2-4ef3-8a0d-8c16b06fd205",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index import (\n",
    "    SimpleDirectoryReader,\n",
    "    VectorStoreIndex,\n",
    "    StorageContext,\n",
    "    load_index_from_storage,\n",
    ")\n",
    "\n",
    "from llama_index.tools import QueryEngineTool, ToolMetadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "91618236-54d3-4783-86b7-7b7554efeed1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    storage_context = StorageContext.from_defaults(persist_dir=\"./storage/lyft\")\n",
    "    lyft_index = load_index_from_storage(storage_context)\n",
    "\n",
    "    storage_context = StorageContext.from_defaults(persist_dir=\"./storage/uber\")\n",
    "    uber_index = load_index_from_storage(storage_context)\n",
    "\n",
    "    index_loaded = True\n",
    "except:\n",
    "    index_loaded = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d3d0bb8c-16c8-4946-a9d8-59528cf3952a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "if not index_loaded:\n",
    "    # load data\n",
    "    lyft_docs = SimpleDirectoryReader(\n",
    "        input_files=[\"../data/10k/lyft_2021.pdf\"]\n",
    "    ).load_data()\n",
    "    uber_docs = SimpleDirectoryReader(\n",
    "        input_files=[\"../data/10k/uber_2021.pdf\"]\n",
    "    ).load_data()\n",
    "\n",
    "    # build index\n",
    "    lyft_index = VectorStoreIndex.from_documents(lyft_docs)\n",
    "    uber_index = VectorStoreIndex.from_documents(uber_docs)\n",
    "\n",
    "    # persist index\n",
    "    lyft_index.storage_context.persist(persist_dir=\"./storage/lyft\")\n",
    "    uber_index.storage_context.persist(persist_dir=\"./storage/uber\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "31892898-a2dc-43c8-812a-3442feb2108d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "lyft_engine = lyft_index.as_query_engine(similarity_top_k=3)\n",
    "uber_engine = uber_index.as_query_engine(similarity_top_k=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f9f3158a-7647-4442-8de1-4db80723b4d2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "query_engine_tools = [\n",
    "    QueryEngineTool(\n",
    "        query_engine=lyft_engine,\n",
    "        metadata=ToolMetadata(\n",
    "            name=\"lyft_10k\",\n",
    "            description=\"Provides information about Lyft financials for year 2021. \"\n",
    "            \"Use a detailed plain text question as input to the tool.\",\n",
    "        ),\n",
    "    ),\n",
    "    QueryEngineTool(\n",
    "        query_engine=uber_engine,\n",
    "        metadata=ToolMetadata(\n",
    "            name=\"uber_10k\",\n",
    "            description=\"Provides information about Uber financials for year 2021. \"\n",
    "            \"Use a detailed plain text question as input to the tool.\",\n",
    "        ),\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "275c01b1-8dce-4216-9203-1e961b7fc313",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Setup OpenAI Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "32f71a46-bdf6-4365-b1f1-e23a0d913a3d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.agent import OpenAIAgent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "ded93297-fee8-4329-bf37-cf77e87621ae",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "agent = OpenAIAgent.from_tools(query_engine_tools, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22716961-11c3-4ac4-82a8-419f787bc36a",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Let's Try It Out!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "42a1bce0-b398-4937-9008-6cee04368ac4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "===== Entering Chat REPL =====\n",
      "Type \"exit\" to exit.\n",
      "\n",
      "=== Calling Function ===\n",
      "Calling function: lyft_10k with args: {\n",
      "  \"input\": \"What was Lyft's revenue growth in 2021?\"\n",
      "}\n",
      "Got output: \n",
      "Lyft's revenue growth in 2021 was 36%.\n",
      "========================\n",
      "=== Calling Function ===\n",
      "Calling function: uber_10k with args: {\n",
      "  \"input\": \"What was Uber's revenue growth in 2021?\"\n",
      "}\n",
      "Got output: \n",
      "Uber's revenue growth in 2021 was 57%.\n",
      "========================\n",
      "Assistant: Lyft's revenue growth in 2021 was 36%, while Uber's revenue growth in 2021 was 57%.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "agent.chat_repl()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
